import os
from PIL import Image

# 输入目录
input_dir = r'E:\CODE\unet_pytorch_lc\img_out\long_200'
txt_dir = r'E:\postgraduate\大论文数据集\分割数据集\IQI'
crop_dir = r'E:\CODE\unet_pytorch_lc\img_out\short_200'
output_dir = r'E:\postgraduate\大论文\seg_train\short还原_200'

# 确保输出目录存在
os.makedirs(output_dir, exist_ok=True)

# 获取所有原始图片文件
image_files = [f for f in os.listdir(input_dir) if f.endswith('.jpg') or f.endswith('.png')]

x = 1
for image_file in image_files:
    # 构建完整的图片、裁剪图片和txt文件路径
    image_path = os.path.join(input_dir, image_file)
    txt_file_path = os.path.join(txt_dir, f'{os.path.splitext(image_file)[0]}.txt')
    crop_image_path = os.path.join(crop_dir, image_file)

    # 检查对应的txt文件和裁剪图片是否存在
    if not os.path.exists(txt_file_path):
        print(f"Warning: {txt_file_path} not found for {image_file}. Skipping.")
        continue
    if not os.path.exists(crop_image_path):
        print(f"Warning: {crop_image_path} not found for {image_file}. Skipping.")
        continue

    # 读取xx1和xx2的值
    with open(txt_file_path, 'r') as f:
        xx1, xx2 = map(float, f.readline().strip().split())

    # 打开原始图片和裁剪图片
    original_image = Image.open(image_path)
    cropped_image = Image.open(crop_image_path)

    # 将裁剪图片粘贴回原始图片
    original_image.paste(cropped_image, (int(xx1), 0))

    # 构造输出文件路径
    output_path = os.path.join(output_dir, image_file)

    # 保存合并后的图片
    original_image.save(output_path)

    print(f"已经完成{x}张图片")
    x = x + 1

print("还原完成！")
